package main

import (
	"bytes"
	"context"
	"crypto/tls"
	"fmt"
	"github.com/lucas-clemente/quic-go"
	"github.com/lucas-clemente/quic-go/quicvarint"
	"log"
	"net/http"
	"sync"
	"time"
	"webtrans/h3"
)

func getHandler() http.Handler {
	mux := http.NewServeMux()
	mux.Handle("/", http.FileServer(http.Dir("web")))
	return mux
}

func main() {

	addr := ":9999"
	certFile := "key/server.cer"
	keyFile := "key/server.key"

	// http2 part
	go func() {
		err := http.ListenAndServeTLS(addr, certFile, keyFile, getHandler())
		if nil != err {
			panic(err)
		}
	}()

	// quic part
	quicConfig := quic.Config{
		Versions:        []quic.VersionNumber{quic.Version1},
		KeepAlive:       true,
		EnableDatagrams: true,
	}
	keyPair, _ := tls.LoadX509KeyPair(certFile, keyFile)
	tlsConfig := tls.Config{
		Certificates: []tls.Certificate{keyPair},
		NextProtos:   []string{"h3"},
	}
	quicListen, _ := quic.ListenAddr(addr, &tlsConfig, &quicConfig)
	for {
		session, acceptErr := quicListen.Accept(context.Background())
		if nil == acceptErr {
			go HandleSession(session)
		} else {
			log.Printf("accept session error %v\n", acceptErr)
		}
	}
}

func HandleSession(sess quic.Session) {
	defer func() {
		// 循环结束,关闭连接
		log.Println("webtransport 连接关闭")
		_ = sess.CloseWithError(0, "")
	}()
	ctx := sess.Context()

	// 服务器先发SettingFrame
	uniStreamFromServer, openUniStreamFromServerErr := sess.OpenUniStreamSync(ctx)
	if nil != openUniStreamFromServerErr {
		log.Printf("open UniStreamFromServer error %v\n", openUniStreamFromServerErr)
		return
	}

	// 发送流类型
	buf := bytes.Buffer{}
	quicvarint.Write(&buf, h3.ControlStreamType)
	_, writeStreamTypeErr := uniStreamFromServer.Write(buf.Bytes())
	if nil != writeStreamTypeErr {
		log.Printf("write StreamType error %v\n", writeStreamTypeErr)
		return
	}

	// 发送SettingFrame
	settingFrameFromServer := h3.SettingFrame{Payload: map[uint64]uint64{
		h3.SettingH3Datagram:         uint64(1),
		h3.SettingEnableWebTransport: uint64(1),
	}}
	_, writeSettingFrameErr := uniStreamFromServer.Write(settingFrameFromServer.ToBytes())
	if nil != writeSettingFrameErr {
		log.Printf("write SettingFrame error %v\n", writeSettingFrameErr)
		return
	}
	log.Println("服务端SettingFrame发送结束")

	uniStreamFromClient, acceptUniStreamFromClient := sess.AcceptUniStream(ctx)
	if nil != acceptUniStreamFromClient {
		log.Printf("accept UniStreamFromClient error %v\n", acceptUniStreamFromClient)
		return
	}
	// 读取流类型
	reader := quicvarint.NewReader(uniStreamFromClient)
	streamType, _ := quicvarint.Read(reader)
	if h3.ControlStreamType == streamType {
		log.Println("收到客户端的ControlStream")
	} else {
		log.Println("客户端发送的uniStream不是ControlStream,关闭")
		return
	}

	// 读取客户端的SettingFrame
	settingFrameFromClient := h3.ParseSettingFrame(uniStreamFromClient)
	log.Println("收到客户端的SettingFrame")
	for key, val := range settingFrameFromClient.Payload {
		fmt.Printf("0x%-10x: 0x%-8x\n", key, val)
	}

	// 等待客户端的bidStream
	bidStreamFromClient, acceptBidStreamFromClientErr := sess.AcceptStream(ctx)
	if nil != acceptBidStreamFromClientErr {
		log.Printf("accept BidStreamFromClient error %v\n", acceptBidStreamFromClientErr)
		return
	}
	log.Println("收到客户端的BidStream")
	// 读取headers
	headersFrame := h3.ParseHeadersFrame(bidStreamFromClient)
	println("收到客户端的headers")
	for k, v := range headersFrame.Header {
		for index := range v {
			fmt.Printf("%-30s:%-10s\n", k, v[index])
		}
	}
	fmt.Printf("%-30s:%-10s\n", "Protocol", headersFrame.Protocol)
	fmt.Printf("%-30s:%-10s\n", "Authority", headersFrame.Authority)
	fmt.Printf("%-30s:%-10s\n", "Path", headersFrame.Path)
	headersFrameFromServer := h3.HeadersFrame{Header: map[string][]string{
		"sec-webtransport-http3-draft": {"draft02"},
	}}
	_, _ = bidStreamFromClient.Write(headersFrameFromServer.ToBytes(200))
	log.Println("服务端headersFrame发送结束")

	wg := &sync.WaitGroup{}
	wg.Add(3)
	time.Sleep(time.Millisecond * 100)
	TestDatagram(sess)
	TestUniStream(sess, ctx)
	TestBidStream(sess, ctx)
	go HandleDatagram(sess, wg)
	go HandleUniStream(sess, wg, ctx)
	go HandleBidStream(sess, wg, ctx)
	wg.Wait()
}

// HandleDatagram 处理DataGram
func HandleDatagram(sess quic.Session, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		message, receiveMessageErr := sess.ReceiveMessage()
		if nil != receiveMessageErr {
			log.Printf("receive message error %v\n", receiveMessageErr)
			break
		} else {
			datagram := h3.ParseDatagram(message)
			msg := datagram.Payload
			log.Printf("receive message %s\n", msg)

			datagramToSend := h3.Datagram{
				StreamId: 0,
				Payload:  "收到datagram数据:" + msg,
			}

			sendMessageErr := sess.SendMessage(datagramToSend.ToBytes())
			if nil != sendMessageErr {
				log.Printf("send message error %v\n", sendMessageErr)
			}
		}
	}
}

// HandleUniStream 处理UniStream
func HandleUniStream(sess quic.Session, wg *sync.WaitGroup, ctx context.Context) {
	defer wg.Done()
	for {
		uniStream, acceptUniStreamErr := sess.AcceptUniStream(ctx)
		if nil != acceptUniStreamErr {
			log.Printf("custom accept UniStream error %v\n", acceptUniStreamErr)
			break
		} else {
			go func(uniStream quic.ReceiveStream) {
				defer func() {
					log.Println("close custom UniStream")
					uniStream.CancelRead(quic.StreamErrorCode(0))
				}()
				// 读取Type
				reader := quicvarint.NewReader(uniStream)
				sType, readTypeErr := quicvarint.Read(reader)
				if nil != readTypeErr || 0x54 != sType {
					log.Printf("custom uniStream type 异常 %v\n", readTypeErr)
					return
				}
				// 读取SessionId
				sessionId, readSessionIdErr := quicvarint.Read(reader)
				if nil != readSessionIdErr {
					log.Printf("read uniStream sessionId error %v\n", readSessionIdErr)
					return
				}
				log.Printf("uniStream sessionId %d\n", sessionId)

				buf := make([]byte, 1024)
				for {
					cnt, readDataErr := reader.Read(buf)
					if nil != readDataErr {
						log.Printf("read uniStream data error %v\n", readDataErr)
						break
					} else {
						if cnt > 0 {
							msg := string(buf[:cnt])
							log.Printf("收到custom uniStream msg %s\n", msg)
							SendWithUniStream(sess, ctx, "收到uni数据:"+msg)
						}
					}
				}
			}(uniStream)
		}
	}
}

// HandleBidStream 处理BidStream
func HandleBidStream(sess quic.Session, wg *sync.WaitGroup, ctx context.Context) {
	defer wg.Done()
	for {
		bidStream, acceptBidStreamErr := sess.AcceptStream(ctx)
		if nil != acceptBidStreamErr {
			log.Printf("custom accept BidStream error %v\n", acceptBidStreamErr)
			break
		} else {
			go func(bidStream quic.Stream) {
				defer func() {
					log.Println("close custom BidStream")
					_ = bidStream.Close()
				}()
				// 读取Type
				reader := quicvarint.NewReader(bidStream)
				sType, readTypeErr := quicvarint.Read(reader)
				if nil != readTypeErr || 0x41 != sType {
					log.Printf("custom bidStream type 异常 %v\n", readTypeErr)
					return
				}
				// 读取SessionId
				sessionId, readSessionIdErr := quicvarint.Read(reader)
				if nil != readSessionIdErr {
					log.Printf("read bidStream sessionId error %v\n", readSessionIdErr)
					return
				}
				log.Printf("bidStream sessionId %d\n", sessionId)

				writer := quicvarint.NewWriter(bidStream)
				buf := make([]byte, 1024)
				for {
					cnt, readDataErr := reader.Read(buf)
					if nil != readDataErr {
						log.Printf("read bidStream data error %v\n", readDataErr)
						break
					} else {
						if cnt > 0 {
							message := buf[:cnt]
							msg := string(message)
							log.Printf("收到custom bidStream msg %s\n", msg)
							_, writeMsgErr := writer.Write([]byte("收到bid数据:" + msg))
							if nil != writeMsgErr {
								log.Printf("custom bidStream write msg error %v\n", writeMsgErr)
							}
						}
					}
				}
			}(bidStream)
		}
	}
}

func TestDatagram(sess quic.Session) {
	datagramToSend := h3.Datagram{
		StreamId: 0,
		Payload:  "来自服务器的Datagram数据",
	}
	writeErr := sess.SendMessage(datagramToSend.ToBytes())
	if nil != writeErr {
		log.Printf("服务器发送Datagram出错 %v", writeErr)
	}
}

func TestUniStream(sess quic.Session, ctx context.Context) {
	streamSync, openUniStreamErr := sess.OpenUniStreamSync(ctx)
	if nil != openUniStreamErr {
		panic(openUniStreamErr)
	}
	defer func() {
		_ = streamSync.Close()
	}()
	writer := quicvarint.NewWriter(streamSync)
	quicvarint.Write(writer, 0x54)
	quicvarint.Write(writer, 0)
	_, writeErr := writer.Write([]byte("来自服务器的单向流数据"))
	if nil != writeErr {
		log.Printf("服务器发送单向流数据出错 %v", writeErr)
	}
}

func TestBidStream(sess quic.Session, ctx context.Context) {
	streamSync, openBidStreamErr := sess.OpenStreamSync(ctx)
	if nil != openBidStreamErr {
		panic(openBidStreamErr)
	}
	writer := quicvarint.NewWriter(streamSync)
	quicvarint.Write(writer, 0x41)
	quicvarint.Write(writer, 0)
	_, writeErr := writer.Write([]byte("来自服务器的双向流数据"))
	if nil != writeErr {
		log.Printf("服务器发送双向流数据出错 %v", writeErr)
	}
}

func SendWithUniStream(sess quic.Session, ctx context.Context, data string) {
	streamSync, openUniStreamErr := sess.OpenUniStreamSync(ctx)
	if nil != openUniStreamErr {
		panic(openUniStreamErr)
	}
	defer func() {
		_ = streamSync.Close()
	}()
	writer := quicvarint.NewWriter(streamSync)
	quicvarint.Write(writer, 0x54)
	quicvarint.Write(writer, 0)
	_, writeErr := writer.Write([]byte(data))
	if nil != writeErr {
		log.Printf("服务器发送单向流数据出错 %v", writeErr)
	}
}
